{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FastKafka tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[FastKafka](https://fastkafka.airt.ai/) is a powerful and easy-to-use Python library for building asynchronous services that interact with Kafka topics. Built on top of [Pydantic](https://docs.pydantic.dev/), [AIOKafka](https://github.com/aio-libs/aiokafka) and [AsyncAPI](https://www.asyncapi.com/), FastKafka simplifies the process of writing producers and consumers for Kafka topics, handling all the parsing, networking, task scheduling and data generation automatically. With FastKafka, you can quickly prototype and develop high-performance Kafka-based services with minimal code, making it an ideal choice for developers looking to streamline their workflow and accelerate their projects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install\n",
    "\n",
    "FastKafka works on macOS, Linux, and most Unix-style operating systems. You can install it with `pip` as usual:\n",
    "\n",
    "```sh\n",
    "pip install fastkafka\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import fastkafka\n",
    "except:\n",
    "    ! pip install fastkafka"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running in Colab\n",
    "\n",
    "You can start this interactive tutorial in Google Colab by clicking the button below:\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/airtai/fastkafka/blob/main/nbs/guides/Guide_00_FastKafka_Demo.ipynb\" target=”_blank”>\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\" />\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing server code\n",
    "\n",
    "Here is an example python script using FastKafka that takes data from a Kafka topic, makes a prediction using a predictive model, and outputs the prediction to another Kafka topic."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing the demo model\n",
    "\n",
    "First we will prepare our model using the Iris dataset so that we can demonstrate the preditions using FastKafka. The following call downloads the dataset and trains the model.\n",
    "\n",
    "We will wrap the model creation into a lifespan of our app so that the model is created just before the app is started."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from contextlib import asynccontextmanager\n",
    "\n",
    "from sklearn.datasets import load_iris\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "from fastkafka import FastKafka\n",
    "\n",
    "ml_models = {}\n",
    "\n",
    "\n",
    "@asynccontextmanager\n",
    "async def lifespan(app: FastKafka):\n",
    "    # Load the ML model\n",
    "    X, y = load_iris(return_X_y=True)\n",
    "    ml_models[\"iris_predictor\"] = LogisticRegression(random_state=0, max_iter=500).fit(\n",
    "        X, y\n",
    "    )\n",
    "    yield\n",
    "    # Clean up the ML models and release the resources\n",
    "    ml_models.clear()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Messages\n",
    "\n",
    "FastKafka uses [Pydantic](https://docs.pydantic.dev/) to parse input JSON-encoded data into Python objects, making it easy to work with structured data in your Kafka-based applications. Pydantic's [`BaseModel`](https://docs.pydantic.dev/usage/models/) class allows you to define messages using a declarative syntax, making it easy to specify the fields and types of your messages.\n",
    "\n",
    "This example defines two message classes for use in a FastKafka application:\n",
    "\n",
    "- The `IrisInputData` class is used to represent input data for a predictive model. It has four fields of type [`NonNegativeFloat`](https://docs.pydantic.dev/usage/types/#constrained-types), which is a subclass of float that only allows non-negative floating point values.\n",
    "\n",
    "- The `IrisPrediction` class is used to represent the output of the predictive model. It has a single field `species` of type string representing the predicted species.\n",
    "\n",
    "These message classes will be used to parse and validate incoming data in Kafka consumers and producers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field, NonNegativeFloat\n",
    "\n",
    "\n",
    "class IrisInputData(BaseModel):\n",
    "    sepal_length: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Sepal length in cm\"\n",
    "    )\n",
    "    sepal_width: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Sepal width in cm\"\n",
    "    )\n",
    "    petal_length: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Petal length in cm\"\n",
    "    )\n",
    "    petal_width: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Petal width in cm\"\n",
    "    )\n",
    "\n",
    "\n",
    "class IrisPrediction(BaseModel):\n",
    "    species: str = Field(..., example=\"setosa\", description=\"Predicted species\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Application\n",
    "\n",
    "This example shows how to initialize a FastKafka application.\n",
    "\n",
    "It starts by defining  a dictionary called `kafka_brokers`, which contains two entries: `\"localhost\"` and `\"production\"`, specifying local development and production Kafka brokers. Each entry specifies the URL, port, and other details of a Kafka broker. This dictionary is used for generating the documentation only and it is not being checked by the actual server.\n",
    "\n",
    "Next, an object of the `FastKafka` class is initialized with the minimum set of arguments:\n",
    "\n",
    "- `kafka_brokers`: a dictionary used for generation of documentation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastkafka import FastKafka\n",
    "\n",
    "kafka_brokers = {\n",
    "    \"localhost\": {\n",
    "        \"url\": \"localhost\",\n",
    "        \"description\": \"local development kafka broker\",\n",
    "        \"port\": 9092,\n",
    "    },\n",
    "    \"production\": {\n",
    "        \"url\": \"kafka.airt.ai\",\n",
    "        \"description\": \"production kafka broker\",\n",
    "        \"port\": 9092,\n",
    "        \"protocol\": \"kafka-secure\",\n",
    "        \"security\": {\"type\": \"plain\"},\n",
    "    },\n",
    "}\n",
    "\n",
    "kafka_app = FastKafka(\n",
    "    title=\"Iris predictions\",\n",
    "    kafka_brokers=kafka_brokers,\n",
    "    lifespan=lifespan,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function decorators\n",
    "\n",
    "FastKafka provides convenient function decorators `@kafka_app.consumes` and `@kafka_app.produces` to allow you to delegate the actual process of\n",
    "\n",
    "- consuming and producing data to Kafka, and\n",
    "\n",
    "- decoding and encoding JSON encode messages\n",
    "\n",
    "from user defined functions to the framework. The FastKafka framework delegates these jobs to AIOKafka and Pydantic libraries.\n",
    "\n",
    "These decorators make it easy to specify the processing logic for your Kafka consumers and producers, allowing you to focus on the core business logic of your application without worrying about the underlying Kafka integration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This following example shows how to use the `@kafka_app.consumes` and `@kafka_app.produces` decorators in a FastKafka application:\n",
    "\n",
    "- The `@kafka_app.consumes` decorator is applied to the `on_input_data` function, which specifies that this function should be called whenever a message is received on the \"input_data\" Kafka topic. The `on_input_data` function takes a single argument which is expected to be an instance of the `IrisInputData` message class. Specifying the type of the single argument is instructing the Pydantic to use `IrisInputData.parse_raw()` on the consumed message before passing it to the user defined function `on_input_data`.\n",
    "\n",
    "- The `@produces` decorator is applied to the `to_predictions` function, which specifies that this function should produce a message to the \"predictions\" Kafka topic whenever it is called. The `to_predictions` function takes a single integer argument `species_class` representing one of three possible strign values predicted by the mdoel. It creates a new `IrisPrediction` message using this value and then returns it. The framework will call the `IrisPrediction.json().encode(\"utf-8\")` function on the returned value and produce it to the specified topic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@kafka_app.consumes(topic=\"input_data\", auto_offset_reset=\"latest\")\n",
    "async def on_input_data(msg: IrisInputData):\n",
    "    species_class = ml_models[\"iris_predictor\"].predict(\n",
    "        [[msg.sepal_length, msg.sepal_width, msg.petal_length, msg.petal_width]]\n",
    "    )[0]\n",
    "\n",
    "    await to_predictions(species_class)\n",
    "\n",
    "\n",
    "@kafka_app.produces(topic=\"predictions\")\n",
    "async def to_predictions(species_class: int) -> IrisPrediction:\n",
    "    iris_species = [\"setosa\", \"versicolor\", \"virginica\"]\n",
    "\n",
    "    prediction = IrisPrediction(species=iris_species[species_class])\n",
    "    return prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing the service\n",
    "\n",
    "The service can be tested using the `Tester` instances which internally starts Kafka broker and zookeeper.\n",
    "\n",
    "Before running tests, we have to install Java runtime and Apache Kafka locally. To simplify the process, we provide the following convenience command:\n",
    "\n",
    "```sh\n",
    "fastkafka testing install_deps\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\r\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\r\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\r\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\r\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "! fastkafka testing install_deps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\r\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\r\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\r\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\r\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "! fastkafka testing install_deps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._start() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'bootstrap_servers': 'localhost:9092', 'auto_offset_reset': 'latest', 'max_poll_records': 100}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['input_data']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'bootstrap_servers': 'localhost:9092', 'auto_offset_reset': 'earliest', 'max_poll_records': 100}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: _aiokafka_consumer_loop(): Consumer loop shutting down, waiting for send_stream to drain...\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: _aiokafka_consumer_loop(): Consumer loop shutting down, waiting for send_stream to drain...\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "from fastkafka.testing import Tester\n",
    "\n",
    "msg = IrisInputData(\n",
    "    sepal_length=0.1,\n",
    "    sepal_width=0.2,\n",
    "    petal_length=0.3,\n",
    "    petal_width=0.4,\n",
    ")\n",
    "\n",
    "# Start Tester app and create local Kafka broker for testing\n",
    "async with Tester(kafka_app) as tester:\n",
    "    # Send IrisInputData message to input_data topic\n",
    "    await tester.to_input_data(msg)\n",
    "\n",
    "    # Assert that the kafka_app responded with IrisPrediction in predictions topic\n",
    "    await tester.awaited_mocks.on_predictions.assert_awaited_with(\n",
    "        IrisPrediction(species=\"setosa\"), timeout=2\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recap\n",
    "\n",
    "We have created a Iris classification model and encapulated it into our fastkafka application.\n",
    "The app will consume the IrisInputData from the `input_data` topic and produce the predictions to `predictions` topic.\n",
    "\n",
    "To test the app we have:\n",
    "\n",
    "1. Created the app\n",
    "\n",
    "2. Started our Tester class which mirrors the developed app topics for testing purpuoses\n",
    "\n",
    "3. Sent IrisInputData message to `input_data` topic\n",
    "\n",
    "4. Asserted and checked that the developed iris classification service has reacted to IrisInputData message "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the service\n",
    "\n",
    "The service can be started using builtin `faskafka run` CLI command. Before we can do that, we will concatenate the code snippets from above and save them in a file `\"application.py\"`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "```python\n",
       "# content of the \"application.py\" file\n",
       "\n",
       "from contextlib import asynccontextmanager\n",
       "\n",
       "from sklearn.datasets import load_iris\n",
       "from sklearn.linear_model import LogisticRegression\n",
       "\n",
       "from fastkafka import FastKafka\n",
       "\n",
       "ml_models = {}\n",
       "\n",
       "\n",
       "@asynccontextmanager\n",
       "async def lifespan(app: FastKafka):\n",
       "    # Load the ML model\n",
       "    X, y = load_iris(return_X_y=True)\n",
       "    ml_models[\"iris_predictor\"] = LogisticRegression(random_state=0, max_iter=500).fit(\n",
       "        X, y\n",
       "    )\n",
       "    yield\n",
       "    # Clean up the ML models and release the resources\n",
       "    ml_models.clear()\n",
       "\n",
       "\n",
       "from pydantic import BaseModel, NonNegativeFloat, Field\n",
       "\n",
       "class IrisInputData(BaseModel):\n",
       "    sepal_length: NonNegativeFloat = Field(\n",
       "        ..., example=0.5, description=\"Sepal length in cm\"\n",
       "    )\n",
       "    sepal_width: NonNegativeFloat = Field(\n",
       "        ..., example=0.5, description=\"Sepal width in cm\"\n",
       "    )\n",
       "    petal_length: NonNegativeFloat = Field(\n",
       "        ..., example=0.5, description=\"Petal length in cm\"\n",
       "    )\n",
       "    petal_width: NonNegativeFloat = Field(\n",
       "        ..., example=0.5, description=\"Petal width in cm\"\n",
       "    )\n",
       "\n",
       "\n",
       "class IrisPrediction(BaseModel):\n",
       "    species: str = Field(..., example=\"setosa\", description=\"Predicted species\")\n",
       "    \n",
       "from fastkafka import FastKafka\n",
       "\n",
       "kafka_brokers = {\n",
       "    \"localhost\": {\n",
       "        \"url\": \"localhost\",\n",
       "        \"description\": \"local development kafka broker\",\n",
       "        \"port\": 9092,\n",
       "    },\n",
       "    \"production\": {\n",
       "        \"url\": \"kafka.airt.ai\",\n",
       "        \"description\": \"production kafka broker\",\n",
       "        \"port\": 9092,\n",
       "        \"protocol\": \"kafka-secure\",\n",
       "        \"security\": {\"type\": \"plain\"},\n",
       "    },\n",
       "}\n",
       "\n",
       "kafka_app = FastKafka(\n",
       "    title=\"Iris predictions\",\n",
       "    kafka_brokers=kafka_brokers,\n",
       "    lifespan=lifespan,\n",
       ")\n",
       "\n",
       "@kafka_app.consumes(topic=\"input_data\", auto_offset_reset=\"latest\")\n",
       "async def on_input_data(msg: IrisInputData):\n",
       "    species_class = ml_models[\"iris_predictor\"].predict(\n",
       "        [[msg.sepal_length, msg.sepal_width, msg.petal_length, msg.petal_width]]\n",
       "    )[0]\n",
       "\n",
       "    to_predictions(species_class)\n",
       "\n",
       "\n",
       "@kafka_app.produces(topic=\"predictions\")\n",
       "def to_predictions(species_class: int) -> IrisPrediction:\n",
       "    iris_species = [\"setosa\", \"versicolor\", \"virginica\"]\n",
       "\n",
       "    prediction = IrisPrediction(species=iris_species[species_class])\n",
       "    return prediction\n",
       "\n",
       "```\n"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# | echo: false\n",
    "\n",
    "from IPython.display import Markdown\n",
    "\n",
    "kafka_app_source = \"\"\"\n",
    "from contextlib import asynccontextmanager\n",
    "\n",
    "from sklearn.datasets import load_iris\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "from fastkafka import FastKafka\n",
    "\n",
    "ml_models = {}\n",
    "\n",
    "\n",
    "@asynccontextmanager\n",
    "async def lifespan(app: FastKafka):\n",
    "    # Load the ML model\n",
    "    X, y = load_iris(return_X_y=True)\n",
    "    ml_models[\"iris_predictor\"] = LogisticRegression(random_state=0, max_iter=500).fit(\n",
    "        X, y\n",
    "    )\n",
    "    yield\n",
    "    # Clean up the ML models and release the resources\n",
    "    ml_models.clear()\n",
    "\n",
    "\n",
    "from pydantic import BaseModel, NonNegativeFloat, Field\n",
    "\n",
    "class IrisInputData(BaseModel):\n",
    "    sepal_length: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Sepal length in cm\"\n",
    "    )\n",
    "    sepal_width: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Sepal width in cm\"\n",
    "    )\n",
    "    petal_length: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Petal length in cm\"\n",
    "    )\n",
    "    petal_width: NonNegativeFloat = Field(\n",
    "        ..., example=0.5, description=\"Petal width in cm\"\n",
    "    )\n",
    "\n",
    "\n",
    "class IrisPrediction(BaseModel):\n",
    "    species: str = Field(..., example=\"setosa\", description=\"Predicted species\")\n",
    "    \n",
    "from fastkafka import FastKafka\n",
    "\n",
    "kafka_brokers = {\n",
    "    \"localhost\": {\n",
    "        \"url\": \"localhost\",\n",
    "        \"description\": \"local development kafka broker\",\n",
    "        \"port\": 9092,\n",
    "    },\n",
    "    \"production\": {\n",
    "        \"url\": \"kafka.airt.ai\",\n",
    "        \"description\": \"production kafka broker\",\n",
    "        \"port\": 9092,\n",
    "        \"protocol\": \"kafka-secure\",\n",
    "        \"security\": {\"type\": \"plain\"},\n",
    "    },\n",
    "}\n",
    "\n",
    "kafka_app = FastKafka(\n",
    "    title=\"Iris predictions\",\n",
    "    kafka_brokers=kafka_brokers,\n",
    "    lifespan=lifespan,\n",
    ")\n",
    "\n",
    "@kafka_app.consumes(topic=\"input_data\", auto_offset_reset=\"latest\")\n",
    "async def on_input_data(msg: IrisInputData):\n",
    "    species_class = ml_models[\"iris_predictor\"].predict(\n",
    "        [[msg.sepal_length, msg.sepal_width, msg.petal_length, msg.petal_width]]\n",
    "    )[0]\n",
    "\n",
    "    await to_predictions(species_class)\n",
    "\n",
    "\n",
    "@kafka_app.produces(topic=\"predictions\")\n",
    "async def to_predictions(species_class: int) -> IrisPrediction:\n",
    "    iris_species = [\"setosa\", \"versicolor\", \"virginica\"]\n",
    "\n",
    "    prediction = IrisPrediction(species=iris_species[species_class])\n",
    "    return prediction\n",
    "\"\"\"\n",
    "\n",
    "with open(\"application.py\", \"w\") as source:\n",
    "    source.write(kafka_app_source)\n",
    "\n",
    "Markdown(\n",
    "    f\"\"\"\n",
    "```python\n",
    "# content of the \"application.py\" file\n",
    "{kafka_app_source}\n",
    "```\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._start() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker starting\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'bootstrap_servers': 'localhost:9092', 'auto_offset_reset': 'latest', 'max_poll_records': 100}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['input_data']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'bootstrap_servers': 'localhost:9092', 'auto_offset_reset': 'earliest', 'max_poll_records': 100}\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched start() called()\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched subscribe() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer.subscribe(), subscribing to: ['predictions']\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: _aiokafka_consumer_loop(): Consumer loop shutting down, waiting for send_stream to drain...\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: _aiokafka_consumer_loop(): Consumer loop shutting down, waiting for send_stream to drain...\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaConsumer patched stop() called\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[INFO] fastkafka._testing.in_memory_broker: AIOKafkaProducer patched stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker._stop() called\n",
      "[INFO] fastkafka._testing.in_memory_broker: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "# | hide\n",
    "\n",
    "from application import kafka_app\n",
    "\n",
    "from fastkafka.testing import Tester\n",
    "\n",
    "msg = IrisInputData(\n",
    "    sepal_length=0.1,\n",
    "    sepal_width=0.2,\n",
    "    petal_length=0.3,\n",
    "    petal_width=0.4,\n",
    ")\n",
    "\n",
    "# Start Tester app and create InMemory Kafka broker for testing - application.py\n",
    "async with Tester(kafka_app) as tester:\n",
    "    # Send IrisInputData message to input_data topic - application.py\n",
    "    await tester.to_input_data(msg)\n",
    "\n",
    "    # Assert that the kafka_app responded with IrisPrediction in predictions topic - application.py\n",
    "    await tester.awaited_mocks.on_predictions.assert_awaited_with(\n",
    "        IrisPrediction(species=\"setosa\"), timeout=30\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run the service, you will need a running Kafka broker on localhost as specified in the `kafka_brokers` parameter above. We can start the Kafka broker locally using the `ApacheKafkaBroker`. Notice that the same happens automatically in the `Tester` as shown above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.start(): entering...\n",
      "[WARNING] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.start(): (<_UnixSelectorEventLoop running=True closed=False debug=False>) is already running!\n",
      "[WARNING] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.start(): calling nest_asyncio.apply()\n",
      "[INFO] fastkafka._components.test_dependencies: Java is already installed.\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\n",
      "[INFO] fastkafka._components.test_dependencies: Kafka is installed.\n",
      "[INFO] fastkafka._components.test_dependencies: But not exported to PATH, exporting...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting zookeeper...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Starting kafka...\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: Local Kafka broker up and running on 127.0.0.1:9092\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: <class 'fastkafka.testing.ApacheKafkaBroker'>.start(): returning 127.0.0.1:9092\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.start(): exited.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'127.0.0.1:9092'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "from fastkafka.testing import ApacheKafkaBroker\n",
    "\n",
    "broker = ApacheKafkaBroker(apply_nest_asyncio=True)\n",
    "\n",
    "broker.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we start the FastKafka service by running the following command in the folder where the `application.py` file is located:\n",
    "\n",
    "```sh\n",
    "fastkafka run --num-workers=2 --kafka-broker localhost application:kafka_app\n",
    "```\n",
    "\n",
    "In the above command, we use `--num-workers` option to specify how many workers to launch and we use `--kafka-broker` option to specify which kafka broker configuration to use from earlier specified `kafka_brokers`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1200656]: [INFO] fastkafka._application.app: set_kafka_broker() : Setting bootstrap_servers value to 'localhost:9092'\n",
      "[1200656]: [INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[1200656]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[1200656]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'bootstrap_servers': 'localhost:9092', 'auto_offset_reset': 'latest', 'max_poll_records': 100}\n",
      "[1200656]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[1200656]: [INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'input_data'})\n",
      "[1200656]: [INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'input_data'}\n",
      "[1200656]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[1200654]: [INFO] fastkafka._application.app: set_kafka_broker() : Setting bootstrap_servers value to 'localhost:9092'\n",
      "[1200654]: [INFO] fastkafka._application.app: _create_producer() : created producer using the config: '{'bootstrap_servers': 'localhost:9092'}'\n",
      "[1200656]: [WARNING] aiokafka.cluster: Topic input_data is not available during auto-create initialization\n",
      "[1200656]: [INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'input_data': 0}. \n",
      "[1200654]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() starting...\n",
      "[1200654]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer created using the following parameters: {'bootstrap_servers': 'localhost:9092', 'auto_offset_reset': 'latest', 'max_poll_records': 100}\n",
      "[1200654]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer started.\n",
      "[1200654]: [INFO] aiokafka.consumer.subscription_state: Updating subscribed topics to: frozenset({'input_data'})\n",
      "[1200654]: [INFO] aiokafka.consumer.consumer: Subscribed to topic(s): {'input_data'}\n",
      "[1200654]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer subscribed.\n",
      "[1200654]: [WARNING] aiokafka.cluster: Topic input_data is not available during auto-create initialization\n",
      "[1200654]: [INFO] aiokafka.consumer.group_coordinator: Metadata for topic has changed from {} to {'input_data': 0}. \n",
      "[1200654]: [ERROR] aiokafka: Unable connect to node with id 0: [Errno 111] Connect call failed ('192.168.112.2', 9092)\n",
      "[1200654]: [ERROR] aiokafka: Unable to update metadata from [0]\n",
      "[1200656]: [ERROR] aiokafka: Unable connect to node with id 0: [Errno 111] Connect call failed ('192.168.112.2', 9092)\n",
      "[1200656]: [ERROR] aiokafka: Unable to update metadata from [0]\n",
      "^C\n",
      "[1200656]: [INFO] fastkafka._components.aiokafka_consumer_loop: _aiokafka_consumer_loop(): Consumer loop shutting down, waiting for send_stream to drain...\n",
      "[1200656]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[1200656]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "[1200654]: [INFO] fastkafka._components.aiokafka_consumer_loop: _aiokafka_consumer_loop(): Consumer loop shutting down, waiting for send_stream to drain...\n",
      "[1200654]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop(): Consumer stopped.\n",
      "[1200654]: [INFO] fastkafka._components.aiokafka_consumer_loop: aiokafka_consumer_loop() finished.\n",
      "Starting process cleanup, this may take a few seconds...\n",
      "[INFO] fastkafka._server: terminate_asyncio_process(): Terminating the process 1200654...\n",
      "[INFO] fastkafka._server: terminate_asyncio_process(): Terminating the process 1200656...\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "!fastkafka run --num-workers=2 --kafka-broker localhost application:kafka_app"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You need to interupt running of the cell above by selecting `Runtime->Interupt execution` on the toolbar above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can stop the local Kafka Broker:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.stop(): entering...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 1200193...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 1200193 was already terminated.\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Terminating the process 1199820...\n",
      "[INFO] fastkafka._components._subprocess: terminate_asyncio_process(): Process 1199820 was already terminated.\n",
      "[INFO] fastkafka._testing.apache_kafka_broker: ApacheKafkaBroker.stop(): exited.\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "broker.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Documentation\n",
    "\n",
    "The kafka app comes with builtin documentation generation using [AsyncApi HTML generator](https://www.asyncapi.com/tools/generator).\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When running in Colab, we need to update Node.js first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "try:\n",
    "    import google.colab\n",
    "\n",
    "    !npm install -g n\n",
    "    !n lts\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to install all dependancies for the generator using the following command line:\n",
    "```sh\n",
    "fastkafka docs install_deps\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.docs_dependencies: AsyncAPI generator installed\r\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "! fastkafka docs install_deps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To generate the documentation programatically you just need to call the folloving command:\n",
    "\n",
    "```sh\n",
    "fastkafka docs generate application:kafka_app\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.asyncapi: New async specifications generated at: '/work/fastkafka/nbs/guides/asyncapi/spec/asyncapi.yml'\n",
      "[INFO] fastkafka._components.asyncapi: Async docs generated at 'asyncapi/docs'\n",
      "[INFO] fastkafka._components.asyncapi: Output of '$ npx -y -p @asyncapi/generator ag asyncapi/spec/asyncapi.yml @asyncapi/html-template -o asyncapi/docs --force-write'\u001b[32m\n",
      "\n",
      "Done! ✨\u001b[0m\n",
      "\u001b[33mCheck out your shiny new generated files at \u001b[0m\u001b[35m/work/fastkafka/nbs/guides/asyncapi/docs\u001b[0m\u001b[33m.\u001b[0m\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "! fastkafka docs generate application:kafka_app"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ". This will generate the *asyncapi* folder in relative path where all your documentation will be saved. You can check out the content of it with:\n",
    "\n",
    "```sh\n",
    "ls -l asyncapi\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 8\r\n",
      "drwxrwxr-x 4 kumaran kumaran 4096 Mar 21 10:09 docs\r\n",
      "drwxrwxr-x 2 kumaran kumaran 4096 Mar 21 10:09 spec\r\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "! ls -l asyncapi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In docs folder you will find the servable static html file of your documentation. This can also be served using our `fastkafka docs serve` CLI command (more on that in our guides).\n",
    "\n",
    "In spec folder you will find a asyncapi.yml file containing the async API specification of your application. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can locally preview the generated documentation by running the following command:\n",
    "\n",
    "```sh\n",
    "fastkafka docs serve application:kafka_app\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] fastkafka._components.asyncapi: New async specifications generated at: '/work/fastkafka/nbs/guides/asyncapi/spec/asyncapi.yml'\n",
      "[INFO] fastkafka._components.asyncapi: Async docs generated at 'asyncapi/docs'\n",
      "[INFO] fastkafka._components.asyncapi: Output of '$ npx -y -p @asyncapi/generator ag asyncapi/spec/asyncapi.yml @asyncapi/html-template -o asyncapi/docs --force-write'\u001b[32m\n",
      "\n",
      "Done! ✨\u001b[0m\n",
      "\u001b[33mCheck out your shiny new generated files at \u001b[0m\u001b[35m/work/fastkafka/nbs/guides/asyncapi/docs\u001b[0m\u001b[33m.\u001b[0m\n",
      "\n",
      "\n",
      "Serving documentation on http://127.0.0.1:8000\u001b[0m\n",
      "^C\n",
      "Interupting serving of documentation and cleaning up...\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# | notest\n",
    "# | echo: false\n",
    "\n",
    "!fastkafka docs serve application:kafka_app"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the parameters passed to the application constructor, we get the documentation bellow:\n",
    "```python\n",
    "from fastkafka import FastKafka\n",
    "\n",
    "kafka_brokers = {\n",
    "    \"localhost\": {\n",
    "        \"url\": \"localhost\",\n",
    "        \"description\": \"local development kafka broker\",\n",
    "        \"port\": 9092,\n",
    "    },\n",
    "    \"production\": {\n",
    "        \"url\": \"kafka.airt.ai\",\n",
    "        \"description\": \"production kafka broker\",\n",
    "        \"port\": 9092,\n",
    "        \"protocol\": \"kafka-secure\",\n",
    "        \"security\": {\"type\": \"plain\"},\n",
    "    },\n",
    "}\n",
    "\n",
    "kafka_app = FastKafka(\n",
    "    title=\"Iris predictions\",\n",
    "    kafka_brokers=kafka_brokers,\n",
    "    bootstrap_servers=\"localhost:9092\",\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Kafka_servers](https://raw.githubusercontent.com/airtai/fastkafka/main/nbs/images/screenshot-kafka-servers.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following documentation snippet are for the consumer as specified in the code above:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "![Kafka_consumer](https://raw.githubusercontent.com/airtai/fastkafka/main/nbs/images/screenshot-kafka-consumer.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following documentation snippet are for the producer as specified in the code above:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "![Kafka_producer](https://raw.githubusercontent.com/airtai/fastkafka/main/nbs/images/screenshot-kafka-producer.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, all messages as defined as subclasses of *BaseModel* are documented as well: "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "![Kafka_![Kafka_servers](https://raw.githubusercontent.com/airtai/fastkafka/main/nbs/images/screenshot-kafka-messages.png)](https://raw.githubusercontent.com/airtai/fastkafka/main/nbs/images/screenshot-kafka-messages.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
